import os

import torch

from .hook import BeforeRunHook, AfterEpochHook


class CheckpointHook(BeforeRunHook, AfterEpochHook):

    def __init__(self, root: str, resume: str = None, interval=1, name: str = "CheckpointHook"):
        super().__init__(name)
        self.interval = interval
        self.root = root
        self.resume = resume

        if not os.path.exists(root):
            os.mkdir(root)

    def before_run(self, container):
        if self.resume is not None:
            state_dict = torch.load(open(os.path.join(self.root, self.resume), 'rb'))
            model = container.model
            model.load_state_dict(state_dict)

    def after_epoch(self, container):
        model = container.model
        state_dict = model.state_dict()
        filename = os.path.join(self.root, "epoch%d.pth" % container.epoch)
        torch.save(state_dict, open(filename, 'wb'))
